# coding: utf-8
# 2021/3/23 @ tongshiwei

import numpy as np
import torch
from tqdm import tqdm
from torch import nn

from .base import Module


class MCDNet(Module):
    """Matrix Factorization Network"""

    def __init__(self, user_num, item_num, latent_dim):
        super(MCDNet, self).__init__()
        self.user_num = user_num
        self.item_num = item_num
        self.latent_dim = latent_dim
        self.user_embedding = nn.Embedding(self.user_num, self.latent_dim)
        self.item_embedding = nn.Embedding(self.item_num, self.latent_dim)
        self.response = nn.Linear(2 * self.latent_dim, 1)

    def forward(self, user_id, item_id):
        user = self.user_embedding(user_id)
        item = self.item_embedding(item_id)
        return torch.squeeze(
            torch.sigmoid(self.response(torch.cat([user, item], dim=-1))), dim=-1
        )
